import numpy as np
import os
import pickle
import matplotlib.pyplot as plt

work_dir = os.getcwd()
pv_calculation_source_path = work_dir + '/result/detailed_soapy_test_result'
pv_calculation_result_path = work_dir + '/result/pv_calculation_for_detailed_soapy_test_result_supplyment'
model_source_path = pv_calculation_source_path + '/perfect label experiment'
model_result_path = pv_calculation_result_path + '/perfect label experiment'

r0_list = [str(round(0.05 + i * 0.01, 2)) for i in range(16)]
mag_list = [str(i) for i in range(8, 19)]


# define function for the processing of data to obtain the descriptive information of it
def calculate_descriptive(data):
    """"
    the output sequence of outcome is mean, std, media, 25_percentile, 75_percentile, max and min
    """
    data_mean = np.mean(data)
    data_std = np.std(data)
    data_media = np.median(data)
    data_25_percentile = np.percentile(data, 25)
    data_75_percentile = np.percentile(data, 75)
    data_max = np.max(data)
    data_min = np.min(data)
    return {'mean': data_mean, 'std': data_std, 'media': data_media, 'P_25': data_25_percentile,
            'P_75': data_75_percentile, 'max': data_max, 'min': data_min}


# firstly calculate pv value under different r0
pv_for_different_r0 = dict()
for j in r0_list:
    work_source_path = model_source_path + '/different_r0/r0=' + j

    pv_for_specific_r0 = dict()
    pv_for_specific_r0['pv_before_with_radius'] = []
    pv_for_specific_r0['pv_after_with_radius'] = []

    # make directory for specific r0
    work_result_path = model_result_path + '/different_r0'
    if not os.path.exists(work_result_path + '/r0=' + j):
        os.makedirs(work_result_path + '/r0=' + j)

    for z in range(1000):
        phase_before = np.load(work_source_path + '/' + str(z) + '/phase_before_radius.npy', allow_pickle=True)
        phase_after = np.load(work_source_path + '/' + str(z) + '/phase_after_radius.npy', allow_pickle=True)

        pv_before_with_radius = np.max(phase_before) - np.min(phase_before)
        pv_after_with_radius = np.max(phase_after) - np.min(phase_after)

        # initiate variable for the storing of information under specific sample
        specific_sample_dirt = dict()
        specific_sample_dirt['pv_before_with_radius'] = pv_before_with_radius
        specific_sample_dirt['pv_after_with_radius'] = pv_after_with_radius

        # make directory for specific sample
        if not os.path.exists(work_result_path + '/r0=' + j + '/' + str(z)):
            os.makedirs(work_result_path + '/r0=' + j + '/' + str(z))

        # save information for specific sample
        with open(work_result_path + '/r0=' + j + '/' + str(z) + '/specific_sample_dirt', 'wb') as file:
            pickle.dump(specific_sample_dirt, file)

        # save information for specific sample in txt format
        with open(work_result_path + '/r0=' + j + '/' + str(z) + '/specific_sample_dirt.txt', 'w') as file:
            file.write(f'specific_sample_dirt = {specific_sample_dirt}')

        pv_for_specific_r0['pv_before_with_radius'].append(pv_before_with_radius)
        pv_for_specific_r0['pv_after_with_radius'].append(pv_after_with_radius)

        # print essential information on screen to tell the advancement of program
        print(f'r0 = {j}, element_index = {z}')

    # make directory for the saving of overall result under specific r0
    if not os.path.exists(model_result_path + '/different_r0/r0=' + j + '/overall_result'):
        os.makedirs(model_result_path + '/different_r0/r0=' + j + '/overall_result')

    specific_r0_overall_result_path = model_result_path + '/different_r0/r0=' + j + '/overall_result'

    # save and plot information after processing
    with open(specific_r0_overall_result_path + '/overall_result_direc', 'wb') as file:
        pickle.dump(pv_for_specific_r0, file)
    plt.subplot(1, 2, 1)
    plt.hist(pv_for_specific_r0['pv_before_with_radius'])
    plt.xlabel('radius')
    plt.ylabel('count')
    plt.title('pv_before_with_radius')
    plt.subplot(1, 2, 2)
    plt.hist(pv_for_specific_r0['pv_after_with_radius'])
    plt.xlabel('radius')
    plt.ylabel('count')
    plt.title('pv_after_with_radius')
    plt.tight_layout()
    plt.savefig(specific_r0_overall_result_path + '/overall_result')
    plt.close('all')

    # store information under specific r0 to relevant directory
    pv_for_different_r0['r0=' + j] = pv_for_specific_r0

# save overall information for all r0
overall_result_path_r0 = model_result_path + '/different_r0/overall_result'
if not os.path.exists(model_result_path + '/different_r0/overall_result'):
    os.makedirs(model_result_path + '/different_r0/overall_result')
with open(overall_result_path_r0 + '/overall', 'wb') as file:
    pickle.dump(pv_for_different_r0, file)
with open(overall_result_path_r0 + '/overall.txt', 'w') as file:
    file.write(f'pv_for_different_r0 = {pv_for_different_r0}')

# prepare and save descriptive information for all r0
# Note!!!!! Instead of re-establish a completely new variable,
# we just modify the variable storing the original information

for j in pv_for_different_r0:

    pv_before_descriptive = calculate_descriptive(pv_for_different_r0[j]['pv_before_with_radius'])
    pv_for_different_r0[j]['pv_before_with_radius'] = pv_before_descriptive

    pv_after_descriptive = calculate_descriptive(pv_for_different_r0[j]['pv_after_with_radius'])
    pv_for_different_r0[j]['pv_after_with_radius'] = pv_after_descriptive

with open(overall_result_path_r0 + '/descriptive_info', 'wb') as file:
    pickle.dump(pv_for_different_r0, file)
with open(overall_result_path_r0 + '/descriptive_info.txt', 'w') as file:
    file.write(f'pv_for_different_r0 = {pv_for_different_r0}')

# secondly calculate pv value under different magnitude
pv_for_different_mag = dict()
for j in mag_list:
    work_source_path = model_source_path + '/different_mag/star_mag=' + j

    pv_for_specific_mag = dict()
    pv_for_specific_mag['pv_before_with_radius'] = []
    pv_for_specific_mag['pv_after_with_radius'] = []

    # make directory for specific magnitude
    work_result_path = model_result_path + '/different_mag'
    if not os.path.exists(work_result_path + '/star_mag=' + j):
        os.makedirs(work_result_path + '/star_mag=' + j)

    for z in range(1000):
        phase_before = np.load(work_source_path + '/' + str(z) + '/phase_before_radius.npy', allow_pickle=True)
        phase_after = np.load(work_source_path + '/' + str(z) + '/phase_after_radius.npy', allow_pickle=True)

        pv_before_with_radius = np.max(phase_before) - np.min(phase_before)
        pv_after_with_radius = np.max(phase_after) - np.min(phase_after)

        # initiate variable for the storing of information under specific sample
        specific_sample_dirt = dict()
        specific_sample_dirt['pv_before_with_radius'] = pv_before_with_radius
        specific_sample_dirt['pv_after_with_radius'] = pv_after_with_radius

        # make directory for specific sample
        if not os.path.exists(work_result_path + '/star_mag=' + j + '/' + str(z)):
            os.makedirs(work_result_path + '/star_mag=' + j + '/' + str(z))

        # save information for specific sample
        with open(work_result_path + '/star_mag=' + j + '/' + str(z) + '/specific_sample_dirt', 'wb') as file:
            pickle.dump(specific_sample_dirt, file)

        # save information for specific sample in txt format
        with open(work_result_path + '/star_mag=' + j + '/' + str(z) + '/specific_sample_dirt.txt', 'w') as file:
            file.write(f'specific_sample_dirt = {specific_sample_dirt}')

        pv_for_specific_mag['pv_before_with_radius'].append(pv_before_with_radius)
        pv_for_specific_mag['pv_after_with_radius'].append(pv_after_with_radius)

        # print essential information on screen to tell the advancement of program
        print(f'magnitude = {j}, element_index = {z}')

    # make directory for the saving of overall result under specific magnitude
    if not os.path.exists(model_result_path + '/different_mag/star_mag=' + j + '/overall_result'):
        os.makedirs(model_result_path + '/different_mag/star_mag=' + j + '/overall_result')

    specific_mag_overall_result_path = model_result_path + '/different_mag/star_mag=' + j + '/overall_result'

    # save and plot information after processing
    with open(specific_mag_overall_result_path + '/overall_result_direc', 'wb') as file:
        pickle.dump(pv_for_specific_mag, file)
    plt.subplot(1, 2, 1)
    plt.hist(pv_for_specific_mag['pv_before_with_radius'])
    plt.xlabel('radius')
    plt.ylabel('count')
    plt.title('pv_before_with_radius')
    plt.subplot(1, 2, 2)
    plt.hist(pv_for_specific_mag['pv_after_with_radius'])
    plt.xlabel('radius')
    plt.ylabel('count')
    plt.title('pv_after_with_radius')
    plt.tight_layout()
    plt.savefig(specific_mag_overall_result_path + '/overall_result')
    plt.close('all')

    # store information under specific magnitude to relevant directory
    pv_for_different_mag['star_mag=' + j] = pv_for_specific_mag

# save overall information for all magnitude
overall_result_path_mag = model_result_path + '/different_mag/overall_result'
if not os.path.exists(model_result_path + '/different_mag/overall_result'):
    os.makedirs(model_result_path + '/different_mag/overall_result')
with open(overall_result_path_mag + '/overall', 'wb') as file:
    pickle.dump(pv_for_different_mag, file)
with open(overall_result_path_mag + '/overall.txt', 'w') as file:
    file.write(f'pv_for_different_mag = {pv_for_different_mag}')

# prepare and save descriptive information for all magnitude
# Note!!!!! Instead of re-establish a completely new variable,
# we just modify the variable storing the original information

for j in pv_for_different_mag:
    pv_before_descriptive = calculate_descriptive(pv_for_different_mag[j]['pv_before_with_radius'])
    pv_for_different_mag[j]['pv_before_with_radius'] = pv_before_descriptive

    pv_after_descriptive = calculate_descriptive(pv_for_different_mag[j]['pv_after_with_radius'])
    pv_for_different_mag[j]['pv_after_with_radius'] = pv_after_descriptive

with open(overall_result_path_mag + '/descriptive_info', 'wb') as file:
    pickle.dump(pv_for_different_mag, file)
with open(overall_result_path_mag + '/descriptive_info.txt', 'w') as file:
    file.write(f'pv_for_different_mag = {pv_for_different_mag}')
